import pandas as pd


def process_combinations(df):
    # 将日期列转换为日期时间格式
    df['trading_date'] = pd.to_datetime(df['trading_date'])

    # 先按comb_id和trading_date进行排序
    df = df.sort_values(by=['comb_id', 'trading_date'])

    # 初始化 start_date 和 end_date 列
    df['start_date'] = df['trading_date']
    df['end_date'] = df['trading_date']

    # 遍历每个组合
    for comb_id, comb_group in df.groupby('comb_id'):
        # 获取该组合的第一个 trading_date
        first_trading_date = comb_group['trading_date'].iloc[0]
        prev_date = None  # 前一个日期
        prev_index = None  # 前一个索引

        update_flag = False
        # 遍历组合内的每一行
        for index, row in comb_group.iterrows():
            # 如果prev_date不为空，即不是第一行

            if prev_date is not None:
                # 当前日期和前一个日期的差值
                date_diff = (row['trading_date'] - prev_date).days

                # 如果差值大于等于10，更新start_date为当前日期
                if date_diff >= 10:
                    df.loc[index, 'start_date'] = row['trading_date']
                    update_flag = False
                    prev_date = None

            # 如果股票组合相关性第一次大于0.8，则开始更新start_date, end_date
            if not(update_flag) and row['corr'] >= 0.8:
                first_trading_date = row['trading_date']
                update_flag = True

            if update_flag:

                # 否则更新end_date为下一行的日期
                df.loc[index, 'start_date'] = first_trading_date
                df.loc[index, 'end_date'] = row['trading_date']

                # 更新prev_date和prev_index
                prev_date = row['trading_date']
                prev_index = index

        # 处理最后一行
        df.loc[prev_index, 'start_date'] = first_trading_date
        df.loc[prev_index, 'end_date'] = prev_date
        df['duration'] = (df['end_date'] - df['start_date']).dt.days

    return df

# 测试
# data = {
#     'comb_id': [539261226]*7 + [66602460]*4,
#     'trading_date': ['2022-01-01', '2022-01-02', '2022-02-03', '2022-02-04', '2022-02-05', '2022-01-06', '2022-01-07', '2022-01-01', '2022-01-02', '2022-01-03', '2022-01-04'],
#     'code1': ['x']*11,
#     'code2': ['y']*7 + ['z']*4,
#     'corr': [0.799687, 0.899583, 0.999332, 0.998822, 0.755565, 0.856078, 0.757136, 1.0, 1.0, 1.0, 1.0],
# }

# 测试
data = {
    'comb_id': [539261226]*7,
    'trading_date': ['2022-01-01', '2022-01-02', '2022-02-03', '2022-02-04', '2022-02-05', '2022-01-06', '2022-01-07'],
    'code1': ['x']*7,
    'code2': ['y']*7 ,
    'corr': [0.799687, 0.799583, 0.999332, 0.998822, 0.755565, 0.856078, 0.757136],
}

df = pd.DataFrame(data)
print(df)
result_df = process_combinations(df)
print(result_df.tail(100))

